package middleware

import (
	"context"
	_ "embed"
	"encoding/json"
	"github.com/redis/go-redis/v9"
	"github.com/zeromicro/go-zero/core/logx"
	"github.com/zeromicro/go-zero/rest/httpx"
	"mall-backend/api/internal/config"
	"net/http"
	"strconv"
	"strings"
)

const (
	// 未明确定义的其他限流配置
	LimitOther = "other"
	// 超过限制
	OverQuota = 0
	// 允许
	Allowed = 1
	// 刚好达到限制
	HitQuota = 2
	// 达到告警值
	HitWarn = 3
)

var (
	//go:embed periodscript.lua
	periodLuaScript string
)

// 限流中间件
type LimitMiddleWare struct {
	cfg          config.Config
	store        *redis.ClusterClient
	periodScript *redis.Script
	excludes     map[string]struct{}
	items        map[string]config.FlowLimitItem
}

func NewLimitMiddleWare(c config.Config, redisCluster *redis.ClusterClient) *LimitMiddleWare {
	middleware := &LimitMiddleWare{
		cfg:      c,
		store:    redisCluster,
		excludes: make(map[string]struct{}),
		items:    make(map[string]config.FlowLimitItem),
	}
	for _, path := range c.FlowLimit.Excludes {
		middleware.excludes[path] = struct{}{}
	}
	for _, item := range c.FlowLimit.Items {
		item = item
		if item.Warn == 0 || item.Warn > item.Quota {
			item.Warn = item.Quota
		}
		middleware.items[item.Path] = item
	}
	middleware.periodScript = redis.NewScript(periodLuaScript)
	return middleware
}

// Handle 限流处理
func (m *LimitMiddleWare) Handle(next http.HandlerFunc) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		path := r.URL.Path
		// 限流关闭，放行
		if !m.cfg.FlowLimit.Enable {
			next(w, r)
			return
		}
		// 此接口不限流，直接放行
		if _, ok := m.excludes[path]; ok {
			next(w, r)
			return
		}
		// 接口单独定义了限流
		if limitItem, ok := m.items[path]; ok {
			// 接口定义的限流关闭，放行
			if !limitItem.Enable {
				next(w, r)
				return
			}
			m.takeCtx(r.Context(), limitItem, next, w, r)
			return
		}
		// 定义了其他总限流
		if limitItem, ok := m.items[LimitOther]; ok {
			// 其他限流关闭，放行
			if !limitItem.Enable {
				next(w, r)
				return
			}
			m.takeCtx(r.Context(), limitItem, next, w, r)
			return
		}
		// 没有限流， 直接放行
		next(w, r)
	}
}

// redis限流
func (m *LimitMiddleWare) takeCtx(ctx context.Context, item config.FlowLimitItem, next http.HandlerFunc, w http.ResponseWriter, r *http.Request) int64 {
	args := []string{}
	for _, ele := range item.Elements {
		val := ctx.Value(ele)
		if val == nil {
			continue
		}
		switch ele {
		case "device_id", "client_ip":
			args = append(args, val.(string))
		case "uid":
			uid, err := val.(json.Number).Int64()
			if err != nil {
				continue
			}
			args = append(args, strconv.FormatInt(uid, 10))
		default:
			continue
		}
	}
	key := item.Prefix + strings.Join(args, ":")
	resp, err := m.periodScript.Run(ctx, m.store, []string{key}, []string{
		strconv.Itoa(item.Quota),
		strconv.Itoa(item.Period),
		strconv.Itoa(item.Warn),
	}).Result()
	if err != nil {
		logx.Error("TakeCtx redis lua script run err:", err)
		return OverQuota
	}
	code, ok := resp.(int64)
	if !ok {
		logx.Error("TakeCtx redis lua script result not int64")
		return OverQuota
	}
	switch code {
	//放行
	case Allowed:
		next(w, r)
	//达到限制值
	case HitQuota:
		logx.Infof("TakeCtx limit key:%s path:%s hit quota", key, r.URL.Path)
		next(w, r)
	//达到告警值
	case HitWarn:
		logx.Infof("TakeCtx limit key:%s path:%s hit warn", key, r.URL.Path)
		next(w, r)
	//触发限流，拒绝请求
	default:
		logx.Errorf("TakeCtx limit key:%s path:%s over quota", key, r.URL.Path)
		httpx.WriteJson(w, http.StatusBadRequest, map[string]interface{}{
			"code":    -1,
			"message": "request over limit",
		})
	}
	return code
}
